/**************************************************************************/
/*                                                                        */
/*       Copyright (c) Microsoft Corporation. All rights reserved.        */
/*                                                                        */
/*       This software is licensed under the Microsoft Software License   */
/*       Terms for Microsoft Azure RTOS. Full text of the license can be  */
/*       found in the LICENSE file at https://aka.ms/AzureRTOS_EULA       */
/*       and in the root directory of this software.                      */
/*                                                                        */
/**************************************************************************/

/*
 * Copyright (c) 2023, Arm Limited and Contributors. All rights reserved.
 */

#include "sample_azure_iot_credentials.h"
#include "sample_config.h"

#include "nx_api.h"
#include "nx_azure_iot_adu_agent.h"
#include "nx_azure_iot_adu_agent_psa_driver.h"
#include "nx_azure_iot_cert.h"
#include "nx_azure_iot_ciphersuites.h"
#include "nx_azure_iot_hub_client.h"

#include <inttypes.h>
#include <stdio.h>

/*
 * Note: The PnP model ID comes from the following sample in the NetX Duo repository:
 * addons/azure_iot/samples/sample_azure_iot_embedded_sdk_pnp.c
 * When we update NetX Duo, we also need to ensure this ID matches the one used in
 * the above file. It refers to models submitted to http://github.com/Azure/iot-plugandplay-models
 */
#define SAMPLE_PNP_MODEL_ID "dtmi:com:example:Thermostat;4"

static UINT initialize_iothub(NX_AZURE_IOT_HUB_CLIENT *iothub_client_ptr);
static void on_azure_iot_log(az_log_classification classification, UCHAR *msg, UINT msg_len);
static void on_update_receive_change(NX_AZURE_IOT_ADU_AGENT *adu_agent_ptr,
                                     UINT update_state,
                                     UCHAR *provider,
                                     UINT provider_length,
                                     UCHAR *name,
                                     UINT name_length,
                                     UCHAR *version,
                                     UINT version_length);

static NX_SECURE_X509_CERT root_ca_cert = {0};
static NX_SECURE_X509_CERT root_ca_cert_2 = {0};
static NX_SECURE_X509_CERT root_ca_cert_3 = {0};
static UCHAR nx_azure_iot_tls_metadata_buffer[NX_AZURE_IOT_TLS_METADATA_BUFFER_SIZE] = {0};
static ULONG nx_azure_iot_thread_stack[STACK_SIZE_NX_AZURE_IOT / sizeof(ULONG)] = {0};
static NX_AZURE_IOT nx_azure_iot = {0};
static NX_AZURE_IOT_HUB_CLIENT iothub_client = {0};
static NX_AZURE_IOT_ADU_AGENT iothub_adu_agent = {0};

void sample_entry(NX_IP *ip_ptr,
                  NX_PACKET_POOL *pool_ptr,
                  NX_DNS *dns_ptr,
                  UINT (*unix_time_callback)(ULONG *unix_time))
{
    nx_azure_iot_log_init(on_azure_iot_log);

    UINT status = nx_azure_iot_create(&nx_azure_iot,
                                      (UCHAR *)"Azure IoT",
                                      ip_ptr,
                                      pool_ptr,
                                      dns_ptr,
                                      nx_azure_iot_thread_stack,
                                      sizeof(nx_azure_iot_thread_stack),
                                      THREAD_PRIORITY_NX_AZURE_IOT,
                                      unix_time_callback);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_create!: error code = 0x%08x\r\n", status);
        return;
    }

    status = nx_secure_x509_certificate_initialize(&root_ca_cert,
                                                   (UCHAR *)_nx_azure_iot_root_cert,
                                                   (USHORT)_nx_azure_iot_root_cert_size,
                                                   NX_NULL,
                                                   0,
                                                   NULL,
                                                   0,
                                                   NX_SECURE_X509_KEY_TYPE_NONE);
    if (status != NX_SECURE_X509_SUCCESS) {
        printf("Failed to initialize ROOT CA certificate!: error code = 0x%08x\r\n", status);
        nx_azure_iot_delete(&nx_azure_iot);
        return;
    }

    status = nx_secure_x509_certificate_initialize(&root_ca_cert_2,
                                                   (UCHAR *)_nx_azure_iot_root_cert_2,
                                                   (USHORT)_nx_azure_iot_root_cert_size_2,
                                                   NX_NULL,
                                                   0,
                                                   NULL,
                                                   0,
                                                   NX_SECURE_X509_KEY_TYPE_NONE);
    if (status != NX_SECURE_X509_SUCCESS) {
        printf("Failed to initialize ROOT CA certificate!: error code = 0x%08x\r\n", status);
        nx_azure_iot_delete(&nx_azure_iot);
        return;
    }

    status = nx_secure_x509_certificate_initialize(&root_ca_cert_3,
                                                   (UCHAR *)_nx_azure_iot_root_cert_3,
                                                   (USHORT)_nx_azure_iot_root_cert_size_3,
                                                   NX_NULL,
                                                   0,
                                                   NULL,
                                                   0,
                                                   NX_SECURE_X509_KEY_TYPE_NONE);
    if (status != NX_SECURE_X509_SUCCESS) {
        printf("Failed to initialize ROOT CA certificate!: error code = 0x%08x\r\n", status);
        nx_azure_iot_delete(&nx_azure_iot);
        return;
    }

    status = initialize_iothub(&iothub_client);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed to initialize iothub client: error code = 0x%08x\r\n", status);
        return;
    }

    status = nx_azure_iot_hub_client_connect(&iothub_client, NX_TRUE, NX_WAIT_FOREVER);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_connect: error code = 0x%08x\r\n", status);
        return;
    }
    printf("Connected to IoT Hub\r\n");

    psa_fwu_component_info_t info_ns;
    psa_status_t psa_status = psa_fwu_query(FWU_COMPONENT_ID_NONSECURE, &info_ns);
    if (psa_status != PSA_SUCCESS) {
        printf("Failed to query non-secure firmware information. Error %" PRId32 "\r\n", psa_status);
        return;
    }
    char sample_version[16];
    int length = snprintf(sample_version,
                          sizeof(sample_version),
                          "%" PRIu16 ".%" PRIu16 ".%" PRIu16,
                          (uint16_t)info_ns.version.major,
                          (uint16_t)info_ns.version.minor,
                          info_ns.version.patch);

    status = nx_azure_iot_adu_agent_start(&iothub_adu_agent,
                                          &iothub_client,
                                          (UCHAR *)SAMPLE_MANUFACTURER,
                                          sizeof(SAMPLE_MANUFACTURER) - 1,
                                          (UCHAR *)SAMPLE_MODEL,
                                          sizeof(SAMPLE_MODEL) - 1,
                                          (UCHAR *)sample_version,
                                          length,
                                          on_update_receive_change,
                                          nx_azure_iot_adu_agent_driver_ns);

    if (status == NX_AZURE_IOT_SUCCESS) {
        printf("Azure Device Update agent started\r\n");
        printf("manufacturer: " SAMPLE_MANUFACTURER ", model: " SAMPLE_MODEL ", version: %s\r\n", sample_version);
    } else {
        printf("Failed to start Azure Device Update agent! error code = 0x%08x\r\n", status);
        return;
    }

    status = nx_azure_iot_hub_client_properties_request(&iothub_client, NX_NO_WAIT);
    if (status == NX_AZURE_IOT_SUCCESS) {
        printf("Checking for update...\r\n");
    } else {
        printf("Failed to check for device update: error code = 0x%08x\r\n", status);
        return;
    }
    /* Nothing else to do in the current thread */
}

static UINT initialize_iothub(NX_AZURE_IOT_HUB_CLIENT *iothub_client_ptr)
{
    UINT status;
    UCHAR *iothub_hostname = (UCHAR *)HOST_NAME;
    UCHAR *iothub_device_id = (UCHAR *)DEVICE_ID;
    UINT iothub_hostname_length = sizeof(HOST_NAME) - 1;
    UINT iothub_device_id_length = sizeof(DEVICE_ID) - 1;

    printf("IoTHub Host Name: %.*s; Device ID: %.*s.\r\n",
           iothub_hostname_length,
           iothub_hostname,
           iothub_device_id_length,
           iothub_device_id);

    status = nx_azure_iot_hub_client_initialize(iothub_client_ptr,
                                                &nx_azure_iot,
                                                iothub_hostname,
                                                iothub_hostname_length,
                                                iothub_device_id,
                                                iothub_device_id_length,
                                                (UCHAR *)MODULE_ID,
                                                sizeof(MODULE_ID) - 1,
                                                _nx_azure_iot_tls_supported_crypto,
                                                _nx_azure_iot_tls_supported_crypto_size,
                                                _nx_azure_iot_tls_ciphersuite_map,
                                                _nx_azure_iot_tls_ciphersuite_map_size,
                                                nx_azure_iot_tls_metadata_buffer,
                                                sizeof(nx_azure_iot_tls_metadata_buffer),
                                                &root_ca_cert);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_initialize!: error code = 0x%08x\r\n", status);
        return (status);
    }

    status = nx_azure_iot_hub_client_model_id_set(
        iothub_client_ptr, (const UCHAR *)SAMPLE_PNP_MODEL_ID, sizeof(SAMPLE_PNP_MODEL_ID) - 1);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_model_id_set!: error code = 0x%08x\r\n", status);
        goto end;
    }

    status = nx_azure_iot_hub_client_trusted_cert_add(iothub_client_ptr, &root_ca_cert_2);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_trusted_cert_add!: error code = 0x%08x\r\n", status);
        goto end;
    }

    status = nx_azure_iot_hub_client_trusted_cert_add(iothub_client_ptr, &root_ca_cert_3);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_trusted_cert_add!: error code = 0x%08x\r\n", status);
        goto end;
    }

    status = nx_azure_iot_hub_client_symmetric_key_set(
        iothub_client_ptr, (UCHAR *)DEVICE_SYMMETRIC_KEY, sizeof(DEVICE_SYMMETRIC_KEY) - 1);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Failed on nx_azure_iot_hub_client_symmetric_key_set!\r\n");
        goto end;
    }

    status = nx_azure_iot_hub_client_properties_enable(iothub_client_ptr);
    if (status != NX_AZURE_IOT_SUCCESS) {
        printf("Properties enable failed!: error code = 0x%08x\r\n", status);
        goto end;
    }

end:

    if (status == NX_AZURE_IOT_SUCCESS) {
        printf("Waiting to connect to IoTHub...\r\n");
    } else {
        nx_azure_iot_hub_client_deinitialize(iothub_client_ptr);
    }

    return (status);
}

static void on_azure_iot_log(az_log_classification classification, UCHAR *msg, UINT msg_len)
{
    if (classification == AZ_LOG_IOT_AZURERTOS) {
        printf("%.*s", msg_len, (CHAR *)msg);
    }
}

static void on_update_receive_change(NX_AZURE_IOT_ADU_AGENT *adu_agent_ptr,
                                     UINT update_state,
                                     UCHAR *provider,
                                     UINT provider_length,
                                     UCHAR *name,
                                     UINT name_length,
                                     UCHAR *version,
                                     UINT version_length)
{
    if (update_state == NX_AZURE_IOT_ADU_AGENT_UPDATE_RECEIVED) {
        printf("Received new update: Provider: %.*s; Name: %.*s, Version: %.*s\r\n",
               provider_length,
               provider,
               name_length,
               name,
               version_length,
               version);

        /* Start to download and install update immediately for testing.  */
        nx_azure_iot_adu_agent_update_download_and_install(adu_agent_ptr);
    } else if (update_state == NX_AZURE_IOT_ADU_AGENT_UPDATE_INSTALLED) {
        /* Start to apply update immediately for testing.  */
        nx_azure_iot_adu_agent_update_apply(adu_agent_ptr);
    }
}
